package com.study.algorithm.pathalgorithm.MIniTree.impl;

import com.study.algorithm.datastruct.graph.model.Edge;
import com.study.algorithm.datastruct.graph.model.Graph;
import com.study.algorithm.datastruct.graph.model.Vertex;
import com.study.algorithm.datastruct.tree.model.Tree;

import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * prime 最小生成树
 * <p>
 * 原理：(取点法）
 * 第一步：取集合中距离最短的一条边，其两端的点被放在被选中的点集合中
 * 第二步：取集合中最短的一条边（
 * 不过这次是有要求的，
 * 1）其两端节点，必须、只能有一个节点在被选中节点集合中
 * 2）或者是，其两端节点都在被选中集合中，但是这两点需在使用已经选中的边不能构成通路的情况下
 * 第三步：重复第二步，直到所有的节点都被选中
 */
public class Prime extends MiniSpanningTreeImpl {
    @Override
    public Tree getMiniSpanningTree(Graph graph) {
        Set<Vertex> vertices = graph.getVertices();
        Set<Vertex> verticesCopy = vertices.stream().collect(Collectors.toSet());

        Set<Edge> edges = sortEdgesDesc(graph.getEdges());
        if (edges == null || edges.size() == 0) {
            return null;
        }


        Set<Edge> selectedEdge = new HashSet<>();
        Set<Vertex> selectedVertices = new HashSet<>();
        Edge miniDistanceEdges = edges.iterator().next();


        addVertexAndEdgeToSelected(miniDistanceEdges, selectedEdge, selectedVertices);
        removeVertexAndEdgeFromOrigin(miniDistanceEdges, edges, verticesCopy);


        while (true) {
            if (verticesCopy.size() <= 0 || edges.size() <= 0) {
                break;
            }
            Edge nextEdge = getNextEdge(selectedVertices, edges, selectedEdge);
            if (nextEdge == null) {
                continue;
            }
            addVertexAndEdgeToSelected(nextEdge, selectedEdge, selectedVertices);
            removeVertexAndEdgeFromOrigin(nextEdge, edges, verticesCopy);
        }

        return generateTreeThroughEdges(vertices, selectedEdge);
    }


    private Edge getNextEdge(Set<Vertex> selectedVertices, Set<Edge> edges, Set<Edge> selectedEdge) {
        Iterator<Edge> it = edges.iterator();
        while (it.hasNext()) {
            Edge edge = it.next();
            Vertex vertex1 = edge.getStartVetex();
            Vertex vertex2 = edge.getEndVetex();
            boolean isSelectedForVertex1 = selectedVertices.contains(vertex1);
            boolean isSelectedForVertex2 = selectedVertices.contains(vertex2);
            if (!isSelectedForVertex1 && !isSelectedForVertex2) {
                continue;
            }
            if ((isSelectedForVertex1 || !isSelectedForVertex2) ||
                    (isSelectedForVertex1 && isSelectedForVertex2)) {
                return edge;
            }
            if (isExistPathFromVertexToVertex(selectedEdge, vertex1, vertex2)) {
                it.remove();
                continue;
            }
            return edge;
        }

        return null;
    }


    private void addVertexAndEdgeToSelected(Edge selectedEdge, Set<Edge> selectedEdges, Set<Vertex> selectedVerties) {
        selectedEdges.add(selectedEdge);
        Vertex vertex1 = selectedEdge.getStartVetex();
        Vertex vertex2 = selectedEdge.getEndVetex();
        selectedVerties.add(vertex1);
        selectedVerties.add(vertex2);
    }

    private void removeVertexAndEdgeFromOrigin(Edge selectedEdge, Set<Edge> edges, Set<Vertex> vertices) {
        Vertex vertex1 = selectedEdge.getStartVetex();
        Vertex vertex2 = selectedEdge.getEndVetex();
        vertices.remove(vertex1);
        vertices.remove(vertex2);
        edges.remove(selectedEdge);
    }

}
